##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(ggplot2)
library(ggpubr)
library(ggsci)
library(cowplot)

##########################################################################################

option_list <- list(
    make_option(c("--class_type"), type = "character") ,
    make_option(c("--input_file"), type = "character") ,
    make_option(c("--base_line_file"), type = "character") ,
    make_option(c("--smg_list"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){
    
    work_dir <- "~/20220915_gastric_multiple/dna_combinePublic"

    class_type <- "DGC"
    input_file <- "~/20220915_gastric_multiple/dna_combinePublic/mutationTime/result/All_CCF_mutTime.tsv"
    base_line_file <- "~/20220915_gastric_multiple/dna_combinePublic/baseTable/STAD_Info.addBurden.MSI_MSS.addCNVType.tsv"

    smg_list <- paste(work_dir,"/mutsig_check/smg.list",sep="")
    out_path <- "~/20220915_gastric_multiple/dna_combinePublic/images/CompareGeneList"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

class_type <- opt$class_type
input_file <- opt$input_file
out_path <- opt$out_path
base_line_file <- opt$base_line_file
smg_list <- opt$smg_list

dir.create( out_path , recursive = T )

###########################################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################
## info
dat_info <- data.frame(fread(base_line_file))

## MutTime
dat_mutTime <- fread(input_file)
#dat_mutTime <- dat_mutTime[which(!(is.na(dat_mutTime$CLS))),]

## mutsig2cv
gene_list <- data.frame(fread(smg_list , header = T))$Gene_Symbol

## 功能性突变
dat_fun <- subset(dat_mutTime , Variant_Classification %in% Variant_Type)

## 驱动突变
dat_driver <- subset(dat_fun , Hugo_Symbol %in% gene_list)

###########################################################################################
## 提取对应类型突变的突变率
dat_info <- subset( dat_info , Class == class_type & Type != "IM + IGC + DGC" )

msi_sample <- unique(subset( dat_info , TCGA_Class=="MSI" )$Patient)
dat_info <- subset( dat_info , !(Patient %in% msi_sample) )

###########################################################################################

dat_driver[grep("In_Frame",dat_driver$Variant_Classification),'Variant_Classification'] = "In_Frame"
dat_driver[grep("Frame_Shift",dat_driver$Variant_Classification),'Variant_Classification'] = "Frame_Shift"
dat_driver <- merge( dat_driver , dat_info[,c("Tumor" , "Patient")] , by.x = "Sample" , by.y = "Tumor" )
## 按照人来计算
dat_driver$Sample <- dat_driver$Patient

###########################################################################################
## 统计突变率
computeRatio <- function(dat_use = dat_use , sample_num = sample_num , type = type){
    dat_mutRate <- dat_use %>%
    group_by( Hugo_Symbol ) %>%
    summarize( MutNum = length(unique(Sample)) )
    dat_mutRate <- data.frame(dat_mutRate)
    dat_mutRate$Ratio <- dat_mutRate$MutNum/sample_num
    colnames( dat_mutRate )[2:3] <- paste0( colnames( dat_mutRate )[2:3] , "_" , type )

    return(dat_mutRate)
}

computeP <- function(dat_combine = dat_combine , sample_num_1 = sample_num_1 , sample_num_2 = sample_num_2){
    result <- data.frame()
    for( gene in  dat_combine$Hugo_Symbol ){

        tmp <- subset( dat_combine , Hugo_Symbol == gene )

        a <- as.numeric(tmp[2])
        c <- as.numeric(sample_num_1 - a)
        b <- as.numeric(tmp[4])
        d <- as.numeric(sample_num_2 - b)

        if(is.na(a)){a=0;c=sample_num_1}
        if(is.na(b)){b=0;d=sample_num_2}

        tmp_result <- fisher.test(matrix(c(a,b,c,d),nrow=2))

        p=tmp_result[["p.value"]]
        OR=round(tmp_result[["estimate"]][["odds ratio"]],3)

        tmp$P <- p
        tmp$OR <- OR

        result <- rbind( result , tmp )
    }

    return(result)
}

hpy_sample <- unique(subset( dat_info , HP == "Positive" )$Patient)
hpn_sample <- unique(subset( dat_info , HP == "Negative" )$Patient)

drink_sample <- unique(subset( dat_info , Alcohol == "Drink" )$Patient)
nondrink_sample <- unique(subset( dat_info , Alcohol == "No" )$Patient)

## HP阳性
dat_use <- subset(dat_driver , Sample %in% hpy_sample )
sample_num <- length(hpy_sample)
type <- "Positive"
dat_mutRate_hpy <- computeRatio(dat_use = dat_use , sample_num = sample_num , type = type)

## HP阴性
dat_use <- subset(dat_driver , Sample %in% hpn_sample )
sample_num <- length(hpn_sample)
type <- "Negative"
dat_mutRate_hpn <- computeRatio(dat_use = dat_use , sample_num = sample_num , type = type)

## 饮酒
use_sample <- drink_sample
type <- "Drink"
dat_use <- subset(dat_driver , Sample %in% use_sample )
sample_num <- length(use_sample)
dat_mutRate_drink <- computeRatio(dat_use = dat_use , sample_num = sample_num , type = type)

## 不饮酒
use_sample <- nondrink_sample
type <- "NonDrink"
dat_use <- subset(dat_driver , Sample %in% use_sample )
sample_num <- length(use_sample)
dat_mutRate_nondrink <- computeRatio(dat_use = dat_use , sample_num = sample_num , type = type)

###########################################################################################
## 合并HP
dat_hp_combine <- merge( dat_mutRate_hpy , dat_mutRate_hpn , all = T )
dat_hp_combine[is.na(dat_hp_combine)] <- 0

## 计算突变的差异
dat_combine <- dat_hp_combine
sample_num_1 <- length(hpy_sample)
sample_num_2 <- length(hpn_sample)

dat_hp_Ratio <- computeP(dat_combine = dat_combine , sample_num_1 = sample_num_1 , sample_num_2 = sample_num_2)

## 合并饮酒
dat_drink_combine <- merge( dat_mutRate_drink , dat_mutRate_nondrink , all = T )
dat_drink_combine[is.na(dat_drink_combine)] <- 0

## 计算突变的差异
dat_combine <- dat_drink_combine
sample_num_1 <- length(drink_sample)
sample_num_2 <- length(nondrink_sample)

dat_drink_Ratio <- computeP(dat_combine = dat_combine , sample_num_1 = sample_num_1 , sample_num_2 = sample_num_2)

###########################################################################################
## 合并最后总的

dat_hpy <- dat_hp_Ratio[,c(1:3,6,7)]
dat_hpy$Type <- "Positive"
dat_hpy$Class <- "HP"
dat_hpn <- dat_hp_Ratio[,c(1,4:7)]
dat_hpn$Type <- "Negative"
dat_hpn$Class <- "HP"

dat_drink <- dat_drink_Ratio[,c(1:3,6,7)]
dat_drink$Type <- "Drinker"
dat_drink$Class <- "Drinking Status"
dat_nondrink <- dat_drink_Ratio[,c(1,4:7)]
dat_nondrink$Type <- "Non-drinker"
dat_nondrink$Class <- "Drinking Status"

col_names <- c( "Hugo_Symbol" , "MutNum" , "Ratio" , "P" , "OR" , "Type" , "Class")
colnames(dat_hpy) <- col_names
colnames(dat_hpn) <- col_names
colnames(dat_drink) <- col_names
colnames(dat_nondrink) <- col_names

result <- rbind( dat_hpy , dat_hpn , dat_drink , dat_nondrink )

###########################################################################################
## gene order
gene_order <- dat_driver %>%
group_by( Hugo_Symbol ) %>%
summarize( MutNum = length(unique(Sample)) )

gene_order <- data.frame(gene_order)
gene_order <- gene_order[order(gene_order$MutNum , decreasing = T),]$Hugo_Symbol

###########################################################################################
## 突变类型计算
computeVariantRatio <- function(dat_driver = dat_driver , sample_use = sample_use , Class = Class , Type = Type){

    dat_driver_tmp <- subset( dat_driver , Sample %in% sample_use )

    dat_driver_tmp_use <- data.frame()
    for(gene in unique(dat_driver_tmp$Hugo_Symbol)){

        tmp <- unique(subset( dat_driver_tmp , Hugo_Symbol == gene )[,c("Hugo_Symbol" , "Sample" , "Variant_Classification")])

        ## 一个人多种突变类型的合并
        tmp_mut <- tmp %>%
        group_by(Sample) %>%
        summarize( MutNum = length(Sample) )

        mutltiSample <- tmp_mut$Sample[tmp_mut$MutNum > 1]
        singleSample <- tmp_mut$Sample[tmp_mut$MutNum == 1]

        if(length(mutltiSample) > 0 ){
            gene_mut_multi <- data.frame( Hugo_Symbol = gene , Sample = mutltiSample , Variant_Classification = "Multiple_Hits" )
        }else{
            gene_mut_multi <- data.frame()
        }
        gene_mut_single <- tmp[tmp$Sample %in% singleSample,c("Hugo_Symbol" , "Sample" , "Variant_Classification")]
        dat_driver_tmp_use <- rbind( dat_driver_tmp_use , gene_mut_multi , gene_mut_single  )
    }


    gene_mut <- dat_driver_tmp_use %>%
    group_by( Hugo_Symbol , Variant_Classification ) %>%
    summarize( MutNum = length(unique(Sample)) )

    gene <- gene_list[!gene_list %in% gene_mut$Hugo_Symbol]
    if(length(gene)>0){
        tmp <- data.frame(Hugo_Symbol = gene , Variant_Classification = "" , MutNum = 0)
        gene_mut <- rbind(gene_mut , tmp)
    }

    gene_mut$Ratio <- gene_mut$MutNum/length(sample_use)
    gene_mut$Class <- Class
    gene_mut$Type <- Type

    return(gene_mut)
}

Class <- "Drinking Status"
sample_use <- nondrink_sample
Type <- "Non-drinker"
gene_mut_nondrink <- computeVariantRatio(dat_driver = dat_driver , sample_use = sample_use , Class = Class , Type = Type)
sample_use <- drink_sample
Type <- "Drinker"
gene_mut_drink <- computeVariantRatio(dat_driver = dat_driver , sample_use = sample_use , Class = Class , Type = Type)

Class <- "HP"
sample_use <- hpy_sample
Type <- "Positive"
gene_mut_hpy <- computeVariantRatio(dat_driver = dat_driver , sample_use = sample_use , Class = Class , Type = Type)
sample_use <- hpn_sample
Type <- "Negative"
gene_mut_hpn <- computeVariantRatio(dat_driver = dat_driver , sample_use = sample_use , Class = Class , Type = Type)

result_gene <- rbind( gene_mut_nondrink , gene_mut_drink , gene_mut_hpn , gene_mut_hpy )

##########################################################################################

result_use <- result
result_use$p_text=ifelse(result_use$P>=0.05,"","*")
result_use$p_text=ifelse(result_use$P<0.05 & result_use$P>0.01,"*",result_use$p_text)
result_use$p_text=ifelse(result_use$P<0.01 & result_use$P>0.001,"**",result_use$p_text)
result_use$p_text=ifelse(result_use$P<0.001 ,"***",result_use$p_text)

result2 <- result_use
result2$value_percent=paste(round(result2$Ratio * 100),"%",sep="")
result2$gene <- factor( result2$Hugo_Symbol , levels = gene_order , order = T )
result2$p_pos <- 0.5
result2$percent_pos <- result2$Ratio + 0.003

##########################################################################################
## 横过来的图
## 分突变类型

col = c(rgb(red=0,green=76,blue=153,alpha=255,max=255),
        rgb(red=191,green=42,blue=55,alpha=255,max=255),
        rgb(red=215,green=114,blue=33,alpha=255,max=255),
        rgb(red=178,green=204,blue=221,alpha=255,max=255) , 
        rgb(red=236,green=130,blue=16,alpha=255,max=255),
        rgb(red=114,green=33,blue=28,alpha=255,max=255)
    )

names(col) = c(
  'Missense_Mutation',
  'Nonsense_Mutation',
  'Frame_Shift',
  'In_Frame',
  'Splice_Site',
  'Multiple_Hits' 
)

col <- col[c(1,4,3,2,5,6)]

result_gene$id <- paste( result_gene$Hugo_Symbol , result_gene$Type , result_gene$Class , sep = ":"  )
result2$id <- paste( result2$Hugo_Symbol , result2$Type , result2$Class , sep = ":"  )

result_gene_combine <- merge( result_gene , 
    result2[,c("id" , "P" , "OR" , "p_text" , "value_percent" , "gene" , "p_pos" , "percent_pos")] , 
    by = "id" )

result_gene_combine_1 <- subset( result_gene_combine , Type == "Positive" )
result_gene_combine_2 <- subset( result_gene_combine , Type == "Negative" )
result_gene_combine_3 <- subset( result_gene_combine , Type == "Drinker" )
result_gene_combine_4 <- subset( result_gene_combine , Type == "Non-drinker" )

result_gene_combine <- rbind( result_gene_combine_1 , result_gene_combine_2 , result_gene_combine_3 , result_gene_combine_4 )
#result_gene_combine <- factor( result_gene_combine$Type , levels = c("Positive" , "Negative" , "Drinker" , "Non-drinker") , order = T)

images_name <- paste0(out_path , "/MutRate.class_type.",class_type,".grid.tsv")
write.table( result_gene_combine , images_name , row.names = F , sep = "\t" , quote = F )

for( class in c("Drinking Status" , "HP") ){

    result3 <- subset( result_gene_combine , Class==class )
    result3$Ratio <- abs(result3$Ratio)
    result3$p_pos <- abs(result3$p_pos)
    result3$percent_pos <- abs(result3$percent_pos)
    result3$gene <- factor( result3$Hugo_Symbol , levels = gene_order[length(gene_order):1] , order = T )
    result3$p_pos <- 0.5

    ## 展示显著差异的基因
    #result4 <- subset( result3 , Ratio > 0.1 )
    result4 <- subset( result3 , P < 0.05 | result3$Ratio == 0 )
    result3 <- subset( result3 , gene %in% unique(result4$gene) )

    for(geneN in unique(result3$gene)){
    
        type <- unique(result3$Type)
        index1 <- which(result3$gene == geneN & result3$Type == type[1])
        index2 <- which(result3$gene == geneN & result3$Type == type[2])
        percent1 <- unique( subset( result3 , gene == geneN & Type == type[1] )$percent_pos)
        percent2 <- unique( subset( result3 , gene == geneN & Type == type[2] )$percent_pos)

        if( percent1 > percent2 ){
            result3[index2,"p_text"] <- ""
        }else{
            result3[index1,"p_text"] <- ""
        }
    }

    plot_bar <- function(result3 = result3 , type = type){

        result <- subset( result3 , Type == type  )
        result$percent_pos <- result$percent_pos + 0.05

        if(type == "Positive" | type == "Drinker" ){
            result$Ratio <- -result$Ratio
            result$p_pos <- -result$p_pos
            result$percent_pos <- -result$percent_pos
        }
        
        result$Variant_Classification <- factor( result$Variant_Classification , levels = names(col)[6:1] , order = T )

        pp <- ggplot(result) + 
        geom_bar(
            aes( x = Ratio , y = gene , fill = Variant_Classification ), color = "white" ,
            stat = "identity",width=0.9)+
        geom_text(aes(label= value_percent , y = gene, x = percent_pos ), 
            position= position_dodge(0.8), vjust=0 , size = 4 , family="Helvetica")+
        geom_text(aes(label= p_text , y = gene , x = p_pos ),size= 5 ,family="Helvetica")+
        scale_fill_manual(values = col) +
        xlab(ifelse(type=="Positive" | type == "Drinker" , unique(result3$Type)[1] , unique(result3$Type)[2]))+
        theme_void()+
        scale_x_continuous(
            breaks = seq(-0.6,0.6,0.2),
            labels = c("60%" , "40%" ,"20%" , "0" , "20%" , "40%" , "60%") ,
            position="top"
        ) +
        #scale_x_discrete(position = "top") +
        theme(
          legend.position = "none",
          axis.title.y = element_blank(),
          panel.border = element_blank(),
          panel.grid=element_blank(),
          panel.grid.major =element_blank(),
          axis.line.y=element_blank(),
          axis.text.y = element_text(size=15),
          axis.text.x = element_text(size=10),
          axis.title.x = element_text(size=15),
          axis.line.x=element_line(linetype=1,color="black",size=0.2)
        )
        if (type == "Positive" | type == "Drinker"){ 
            pp<- pp + 
            theme(axis.text.y =element_blank())
        }
        return(pp)
    }

    type <- unique(result3$Type)[1]
    p1 <- plot_bar(result3 = result3 , type = type)

    type <- unique(result3$Type)[2]
    p2 <- plot_bar(result3 = result3 , type = type)

    pyramid <- plot_grid(p1,p2,ncol=2,align="hv")

    width <- 9
    if(class == "HP"){
        if(class_type == "IGC"){
            height <- 2.2
        }else if(class_type == "DGC"){
            height <- 2.5
        }else if(class_type == "IM"){
            height <- 2
        }
    }

    if(class == "Drinking Status"){
        if(class_type == "IGC"){
            height <- 2.2
        }else if(class_type == "DGC"){
            height <- 2
        }else if(class_type == "IM"){
            height <- 1.8
        }
    }
    

}

images_name <- paste0(out_path , "/MutRate.",class,".",class_type,".grid.pdf")
ggsave( images_name , pyramid , width = width , height = height )